package git.soulbgm.utils;

import org.rocksdb.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * RocksDB 操作工具类
 *
 * @author SoulBGM
 * @date 2022/01/26
 */
public class RocksDBUtil {

    private final static Logger log = LoggerFactory.getLogger(RocksDBUtil.class);
    /**
     * RocksDB
     */
    private static RocksDB rocks;
    /**
     * 数据库列族(表)集合
     */
    public static ConcurrentMap<String, ColumnFamilyHandle> columnFamilyHandleMap = new ConcurrentHashMap<>();

    static {
        String rocksDbPath = System.getProperty("user.dir") + File.separator + "rocksdb";
        RocksDB.loadLibrary();
        Options options = new Options();
        //如果数据库不存在则创建
        options.setCreateIfMissing(true);
        // 初始化所有已存在列族
        try {
            List<byte[]> cfArr = RocksDB.listColumnFamilies(options, rocksDbPath);
            //ColumnFamilyDescriptor集合
            List<ColumnFamilyDescriptor> columnFamilyDescriptors = new ArrayList<>();
            if (cfArr.size() > 0) {
                for (byte[] cf : cfArr) {
                    columnFamilyDescriptors.add(new ColumnFamilyDescriptor(cf, new ColumnFamilyOptions()));
                }
            } else {
                columnFamilyDescriptors.add(new ColumnFamilyDescriptor(RocksDB.DEFAULT_COLUMN_FAMILY, new ColumnFamilyOptions()));
            }
            DBOptions dbOptions = new DBOptions();
            dbOptions.setCreateIfMissing(true);
            //ColumnFamilyHandle集合
            List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>();
            rocks = RocksDB.open(dbOptions, rocksDbPath, columnFamilyDescriptors, columnFamilyHandles);
            for (int i = 0; i < columnFamilyDescriptors.size(); i++) {
                ColumnFamilyHandle columnFamilyHandle = columnFamilyHandles.get(i);
                String cfName = new String(columnFamilyDescriptors.get(i).getName());
                columnFamilyHandleMap.put(cfName, columnFamilyHandle);
            }
            log.info("RocksDB 初始化成功 路径:{}", rocksDbPath);
            log.info("列族有:{}", columnFamilyHandleMap.keySet());
        } catch (RocksDBException e) {
            e.printStackTrace();
        }
    }

    /**
     * 如果不存在列族则进行添加
     *
     * @param cfName 列族名称
     * @return {@link ColumnFamilyHandle}
     * @throws RocksDBException RocksDB异常
     */
    public static ColumnFamilyHandle cfAddIfNotExist(String cfName) throws RocksDBException {
        ColumnFamilyHandle columnFamilyHandle;
        if (!columnFamilyHandleMap.containsKey(cfName)) {
            columnFamilyHandle = rocks.createColumnFamily(new ColumnFamilyDescriptor(cfName.getBytes(), new ColumnFamilyOptions()));
            columnFamilyHandleMap.put(cfName, columnFamilyHandle);
            log.info("列族添加成功 列族名称:{}", cfName);
        } else {
            columnFamilyHandle = columnFamilyHandleMap.get(cfName);
        }
        return columnFamilyHandle;
    }

    /**
     * 如果存在列族则进行删除
     *
     * @param cfName 列族名称
     * @throws RocksDBException RocksDB异常
     */
    public static void cfDeleteIfExist(String cfName) throws RocksDBException {
        if (columnFamilyHandleMap.containsKey(cfName)) {
            rocks.dropColumnFamily(columnFamilyHandleMap.get(cfName));
            columnFamilyHandleMap.remove(cfName);
            log.info("列族删除成功 列族名称:{}", cfName);
        } else {
            log.warn("列族不存在 列族名称:{}", cfName);
        }
    }

    /**
     * 增
     *
     * @param cfName 列族名称
     * @param key    键
     * @param value  值
     * @throws RocksDBException RocksDB异常
     */
    public static void put(String cfName, String key, String value) throws RocksDBException {
        put(cfName, key, value.getBytes());
    }

    /**
     * 增
     *
     * @param cfName 列族名称
     * @param key    键
     * @param value  值
     * @throws RocksDBException RocksDB异常
     */
    public static void put(String cfName, String key, byte[] value) throws RocksDBException {
        //获取列族Handle
        ColumnFamilyHandle columnFamilyHandle = cfAddIfNotExist(cfName);
        rocks.put(columnFamilyHandle, key.getBytes(), value);
    }

    /**
     * 增（批量）
     *
     * @param cfName 列族名称
     * @param map    数据
     * @throws RocksDBException RocksDB异常
     */
    public static void batchPut(String cfName, Map<String, String> map) throws RocksDBException {
        //获取列族Handle
        ColumnFamilyHandle columnFamilyHandle = cfAddIfNotExist(cfName);
        WriteOptions writeOptions = new WriteOptions();
        WriteBatch writeBatch = new WriteBatch();
        for (Map.Entry<String, String> entry : map.entrySet()) {
            writeBatch.put(columnFamilyHandle, entry.getKey().getBytes(), entry.getValue().getBytes());
        }
        rocks.write(writeOptions, writeBatch);
    }

    /**
     * 删
     *
     * @param cfName 列族名称
     * @param key    键
     * @throws RocksDBException RocksDB异常
     */
    public static void delete(String cfName, String key) throws RocksDBException {
        //获取列族Handle
        ColumnFamilyHandle columnFamilyHandle = cfAddIfNotExist(cfName);
        rocks.delete(columnFamilyHandle, key.getBytes());
    }

    /**
     * 查
     *
     * @param cfName 列族名称
     * @param key    键
     * @return {@link byte[]}
     * @throws RocksDBException RocksDB异常
     */
    public static byte[] get(String cfName, String key) throws RocksDBException {
        //获取列族Handle
        ColumnFamilyHandle columnFamilyHandle = cfAddIfNotExist(cfName);
        return rocks.get(columnFamilyHandle, key.getBytes());
    }

    /**
     * 查（多个键值对）
     *
     * @param cfName 列族名称
     * @param keys   键
     * @return {@link Map}<{@link String}, {@link byte[]}>
     * @throws RocksDBException RocksDB异常
     */
    public static Map<String, byte[]> multiGetAsMap(String cfName, String[] keys) throws RocksDBException {
        Map<String, byte[]> map = new HashMap<>();
        //获取列族Handle
        ColumnFamilyHandle columnFamilyHandle = cfAddIfNotExist(cfName);
        List<ColumnFamilyHandle> columnFamilyHandles = new ArrayList<>(keys.length + 1);
        List<byte[]> keyBytes = new ArrayList<>();
        for (String key : keys) {
            keyBytes.add(key.getBytes());
        }
        for (int i = 0; i < keys.length; i++) {
            columnFamilyHandles.add(columnFamilyHandle);
        }
        List<byte[]> bytes = rocks.multiGetAsList(columnFamilyHandles, keyBytes);
        for (int i = 0; i < bytes.size(); i++) {
            byte[] valueBytes = bytes.get(i);
            map.put(keys[i], valueBytes);
        }
        return map;
    }

    /**
     * 查（所有键值对）
     *
     * @param cfName 列族名称
     * @return {@link Map}<{@link String}, {@link byte[]}>
     * @throws RocksDBException RocksDB异常
     */
    public static Map<String, byte[]> getAll(String cfName) throws RocksDBException {
        Map<String, byte[]> map = new HashMap<>();
        //获取列族Handle
        ColumnFamilyHandle columnFamilyHandle = cfAddIfNotExist(cfName);
        RocksIterator rocksIterator = rocks.newIterator(columnFamilyHandle);
        for (rocksIterator.seekToFirst(); rocksIterator.isValid(); rocksIterator.next()) {
            map.put(new String(rocksIterator.key()), rocksIterator.value());
        }
        return map;
    }

}
